import time
import datetime
import unittest
import pandas as pd
from dev.stock_information.stock_detail import get_stock_MA,get_stock_MACD,get_stock_RSI,get_Stock_Bollinger,get_stock_OBV,get_daily_stock
from dev.stock_visual.stock_mpf_plot import KLine_plot,MA_plot,RSI_plot,MACD_plot,Bollinger_plot,OBV_plot
from dev.stock_information.stock_detail import get_aStock_zh_a_minute,get_week_stock,get_month_stock

import sys
import os
from pathlib import Path
curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(curPath)[0]
sys.path.append(rootPath)

stock_code_path = os.path.join(Path(curPath).parent,"resources","2023-02-20_stockCode.csv")

def get_stocks_name(stock_codes_path=stock_code_path):
    res_df = pd.read_csv(stock_codes_path)
    codes = list(res_df['代码'].values)
    names = list(res_df['名称'].values)
    code_name_dict = {}
    for i in range(len(codes)):
        code_name_dict[codes[i]] = names[i]
    return code_name_dict

class StockPlotTest(unittest.TestCase):

    #绘制某只股票的分时数据
    def test_KLine_plot(self):
        codes = ['sh601728']  #中国电信
        period,adjust = "1",""  #获取每1分钟的分时数据
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            df = get_aStock_zh_a_minute(symbol=code, period=period, adjust=adjust)  #分时数据适合按天来分析股票（因此要截取这一天的数据），否则计算的MA对于日K，月K并不能起到平滑作用
            #按时间过滤dataframe 参考https://blog.csdn.net/qq_39898780/article/details/126777300
            start_date_str = (datetime.datetime.now()-datetime.timedelta(days=1)).strftime('%Y-%m-%d %H:%M:%S')
            df = df[df['day'] > start_date_str]

            title = f"{code_name_dict[code]}每{period}分钟的行情数据"
            savePath_temp = tokens[0] + str(index) + "." + tokens[1]

            # 参考 https://blog.csdn.net/WHYbeHERE/article/details/109497750
            df.timestamp = pd.to_datetime(df.day, format='%Y-%m-%d %H:%M:%S')  # 将day字符串类型强转为datetime
            df.index = df.timestamp  # 要求df的index为Datetime类型
            # 对oepn,high,low,close,volume列进行强制类型转换，参考https://mp.weixin.qq.com/s?__biz=MzU5NzcyMzA5MQ==&mid=2247483928&idx=1&sn=e829a48cfb54720fc44c0f8df6a15d3e&chksm=fe4e50a4c939d9b2260abb027cbb70da022e52ec11a4ec7f1146670b1424a847cbcd3e619edf&scene=27
            for column in df.columns[1:]:
                df[column] = pd.to_numeric(df[column])  # str转numeric
            KLine_plot(df, title, savePath_temp)

            # MA5
            savePath_temp = tokens[0] + str(index) + "_" + str(5) + "." + tokens[1]
            MA_data = get_stock_MA(stock_data=df["close"], timeperiod=5, matype=0)
            MA_plot(df, MA_data, title=title, savePath=savePath_temp)
            # MA10
            savePath_temp = tokens[0] + str(index) + "_" + str(10) + "." + tokens[1]
            MA_data = get_stock_MA(stock_data=df["close"], timeperiod=10, matype=0)
            MA_plot(df, MA_data, title=title, savePath=savePath_temp)
            # MA30
            savePath_temp = tokens[0] + str(index) + "_" + str(30) + "." + tokens[1]
            MA_data = get_stock_MA(stock_data=df["close"], timeperiod=30, matype=0)
            MA_plot(df, MA_data, title=title, savePath=savePath_temp)

    # 绘制某只股票的周K/月线
    def test_week_month_Kline_plot(self):
        codes = ['sh601728']  # 中国电信
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            # 根据日时股票数据计算MA
            start_date, end_date = "20220101", "20230215"
            df = get_daily_stock(symbol=code, start_date=start_date, end_date=end_date, adjust="")
            df.index = df.date

            df_week = get_week_stock(df)
            title = f"{code_name_dict[code]}从{start_date}到{end_date}的周K线数据"
            savePath_temp = tokens[0] + str(index) + "." + tokens[1]
            KLine_plot(data=df_week, title=title, savePath=savePath_temp)

            df_month = get_month_stock(df)
            title = f"{code_name_dict[code]}从{start_date}到{end_date}的月K线数据"
            savePath_temp = tokens[0] + str(index) + "_" + str(index) + "." + tokens[1]
            KLine_plot(data=df_month, title=title, savePath=savePath_temp)


    # 绘制某只股票的日时数据
    def test_MA_plot(self):
        codes = ['sh601728']  # 中国电信
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            #根据日时股票数据计算MA
            start_date, end_date = "20220101", "20230215"
            df = get_daily_stock(symbol=code, start_date=start_date,end_date=end_date, adjust="")
            df.index = df.date

            title = f"{code_name_dict[code]}从{start_date}到{end_date}的行情数据"
            savePath_temp = tokens[0] + str(index) +  "." + tokens[1]
            KLine_plot(data=df,title=title,savePath=savePath_temp)

            #MA5
            period = 5
            title = f"{code_name_dict[code]}的MA{period}数据"
            savePath_temp = tokens[0] + str(index) + "_" + str(period) +"." + tokens[1]
            MA_data = get_stock_MA(stock_data=df["close"],timeperiod=period,matype=0)
            MA_plot(df,MA_data,title=title,savePath=savePath_temp)

            # MA5
            period = 10
            title = f"{code_name_dict[code]}的MA{period}数据"
            savePath_temp = tokens[0] + str(index) + "_" + str(period) + "." + tokens[1]
            MA_data = get_stock_MA(stock_data=df["close"], timeperiod=period, matype=0)
            MA_plot(df, MA_data, title=title, savePath=savePath_temp)

            # MA5
            period = 30
            title = f"{code_name_dict[code]}的MA{period}数据"
            savePath_temp = tokens[0] + str(index) + "_" + str(period) + "." + tokens[1]
            MA_data = get_stock_MA(stock_data=df["close"], timeperiod=period, matype=0)
            MA_plot(df, MA_data, title=title, savePath=savePath_temp)

    def test_RSI_plot(self):
        codes = ['sh601728']  # 中国电信
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            # 根据日时股票数据计算MA
            start_date, end_date = "20220101", "20230215"
            df = get_daily_stock(symbol=code, start_date=start_date, end_date=end_date, adjust="")
            df.index = df.date

            title = f"{code_name_dict[code]}从{start_date}到{end_date}的行情数据"
            savePath_temp = tokens[0] + str(index) + "." + tokens[1]

            # RSI
            RSI6 = get_stock_RSI(df['close'],timeperiod=6)
            RSI12 = get_stock_RSI(df['close'],timeperiod=12)
            RSI24 = get_stock_RSI(df['close'],timeperiod=24)
            add_plot = RSI_plot(RSI6,RSI12,RSI24,pos=2)
            KLine_plot(data=df, title=title, savePath=savePath_temp, add_plot=add_plot)

    def test_MACD_plot(self):
        codes = ['sh601728']  # 中国电信
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            # 根据日时股票数据计算MA
            start_date, end_date = "20220101", "20230215"
            df = get_daily_stock(symbol=code, start_date=start_date, end_date=end_date, adjust="")
            df.index = df.date

            title = f"{code_name_dict[code]}从{start_date}到{end_date}的行情数据"
            savePath_temp = tokens[0] + str(index) + "." + tokens[1]

            #macd
            macd,macdsignal,macdhist = get_stock_MACD(stock_data=df['close'],fastperiod=12,slowperiod=26,signalperiod=9)
            add_plot = MACD_plot(macd=macd,macdSignal=macdsignal,histogram=macdhist,pos=2)
            KLine_plot(data=df, title=title, savePath=savePath_temp,add_plot=add_plot)


    def test_Bollinger_plot(self):
        codes = ['sh601728']  # 中国电信
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            # 根据日时股票数据计算MA
            start_date, end_date = "20220101", "20230215"
            df = get_daily_stock(symbol=code, start_date=start_date, end_date=end_date, adjust="")
            df.index = df.date

            title = f"{code_name_dict[code]}从{start_date}到{end_date}的行情数据"
            savePath_temp = tokens[0] + str(index) + "." + tokens[1]

            # Bollinger
            H_line, M_line, L_line = get_Stock_Bollinger(df['close'],timeperiod=5)
            add_plot = Bollinger_plot(H_line, M_line, L_line,pos=2)
            KLine_plot(data=df, title=title, savePath=savePath_temp, add_plot=add_plot)

    def test_OBV_plot(self):
        codes = ['sh601728']  # 中国电信
        savePath = os.path.join(rootPath, "stock_visual/temp", "temp.png")

        code_name_dict = get_stocks_name()
        tokens = savePath.split(".")
        for index, code in enumerate(codes):
            # 根据日时股票数据计算MA
            start_date, end_date = "20220101", "20230215"
            df = get_daily_stock(symbol=code, start_date=start_date, end_date=end_date, adjust="")
            df.index = df.date

            title = f"{code_name_dict[code]}从{start_date}到{end_date}的行情数据"
            savePath_temp = tokens[0] + str(index) + "." + tokens[1]

            # Bollinger
            H_line, M_line, L_line = get_Stock_Bollinger(df['close'], timeperiod=5)
            add_plot = Bollinger_plot(H_line, M_line, L_line, pos=2,ylabel="Bollinger")
            # OBV
            obv = get_stock_OBV(df)
            add_plot1 = OBV_plot(obv,pos=3,ylabel="能量潮")
            add_plot.extend(add_plot1)
            KLine_plot(data=df, title=title, savePath=savePath_temp, add_plot=add_plot)